{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AES_Adv_1\n",
    "\n",
    "In SCA101, we focused only on a single attack location, the SBox, without much explanation or experimentation. In this lab, we'll be trying some other attack locations to better understand why the SBox was used exclusively in SCA101."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Capture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET = 'TINYAES128C'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\"\n",
    "cd ../../firmware/mcu/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"../Helper_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = '../../firmware/mcu/simpleserial-aes/simpleserial-aes-{}.hex'.format(PLATFORM)\n",
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Capture Traces\n",
    "from tqdm import tnrange, trange\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "ktp = cw.ktp.Basic()\n",
    "\n",
    "traces = []\n",
    "N = 200  # Number of traces\n",
    "project = cw.create_project(\"AES_Adv_1.cwp\")\n",
    "\n",
    "for i in trange(N, desc='Capturing traces'):\n",
    "    key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "\n",
    "    trace = cw.capture_trace(scope, target, text, key)\n",
    "    if trace is None:\n",
    "        continue\n",
    "    project.traces.append(trace)\n",
    "\n",
    "#Convert traces to numpy arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.sbox_output\n",
    "attack = cwa.cpa(project, leak_model)\n",
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results.key_guess())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leak_model.key_schedule_rounds(results.key_guess(), 0, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data = cwa.analyzer_plots(results)\n",
    "def byte_to_color(idx):\n",
    "    return hv.Palette.colormaps['Category20'](idx/16.0)\n",
    "\n",
    "import holoviews as hv\n",
    "from holoviews.operation.datashader import datashade, shade, dynspread, rasterize\n",
    "from holoviews.operation import decimate\n",
    "import pandas as pd, numpy as np\n",
    "\n",
    "a = []\n",
    "b = []\n",
    "hv.extension('bokeh')\n",
    "for i in range(0, 16):\n",
    "    data = plot_data.output_vs_time(i)\n",
    "    a.append(np.array(data[1]))\n",
    "    b.append(np.array(data[2]))\n",
    "    b.append(np.array(data[3]))\n",
    "    \n",
    "pda = pd.DataFrame(a).transpose().rename(str, axis='columns')\n",
    "pdb = pd.DataFrame(b).transpose().rename(str, axis='columns')\n",
    "curve = hv.Curve(pdb['0'], \"Sample\").options(color='black')\n",
    "for i in range(1, 16):\n",
    "    curve *= hv.Curve(pdb[str(i)]).options(color='black')\n",
    "    \n",
    "for i in range(0, 16):\n",
    "    curve *= hv.Curve(pda[str(i)]).options(color=byte_to_color(i))\n",
    "decimate(curve.opts(width=900, height=600))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.sbox_output\n",
    "attack = cwa.cpa(project, leak_model)\n",
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leak_model = cwa.leakage_models.plaintext_key_xor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(cwa.leakage_models)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attack = cwa.cpa(project, leak_model)\n",
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At first glance, our attack appears to have partially failed, recovering roughly half the key bytes. This isn't quite true: we've got two equal key guesses for each byte. Remember back to the \"recovering data from a single bit\" slides: the correct key and its inverse will produce inverted results. This means that they're both have the same correlation with the power trace, just that one is negative and one is positive.\n",
    "\n",
    "So we have two possible guesses for each key. Is this enough to stop us from recovering the key? In this case, it's pretty easy to brute force the correct key. In an actual attack, we might not know the ciphertext, making it much harder to bruteforce the key.\n",
    "\n",
    "The XOR operation is also linear, meaning incorrect key guesses have higher correlations than they would with a non-linear operation like the SBox. If we didn't have as nice of a measurement environment, this might make the difference between being able to tell we've got the key and not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data = cwa.analyzer_plots(results)\n",
    "def byte_to_color(idx):\n",
    "    return hv.Palette.colormaps['Category20'](idx/16.0)\n",
    "\n",
    "import holoviews as hv\n",
    "from holoviews.operation.datashader import datashade, shade, dynspread, rasterize\n",
    "from holoviews.operation import decimate\n",
    "import pandas as pd, numpy as np\n",
    "\n",
    "a = []\n",
    "b = []\n",
    "hv.extension('bokeh')\n",
    "for i in range(0, 16):\n",
    "    data = plot_data.output_vs_time(i)\n",
    "    a.append(np.array(data[1]))\n",
    "    b.append(np.array(data[2]))\n",
    "    b.append(np.array(data[3]))\n",
    "    \n",
    "pda = pd.DataFrame(a).transpose().rename(str, axis='columns')\n",
    "pdb = pd.DataFrame(b).transpose().rename(str, axis='columns')\n",
    "curve = hv.Curve(pdb['0'], \"Sample\").options(color='black')\n",
    "for i in range(1, 16):\n",
    "    curve *= hv.Curve(pdb[str(i)]).options(color='black')\n",
    "    \n",
    "for i in range(0, 16):\n",
    "    curve *= hv.Curve(pda[str(i)]).options(color=byte_to_color(i))\n",
    "decimate(curve.opts(width=900, height=600))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.new_model(ShiftRowsOutput)\n",
    "attack = cwa.cpa(project, leak_model)\n",
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the shift rows output, we can see that the attack has again failed. This time, it's much worse: we've recovered only 4 bytes. All the rest have completely failed! You might've noticed something odd about the bytes that failed: every key guess has exactly the same correlation! To understand this, we'll need to look at what the shift rows instruction actuall does. As the name suggests, it shifts the rows of the AES state around, more specifically, it turns the following matrix:\n",
    "\n",
    "$$\n",
    "state = \\left( \\begin{array}\n",
    "& S0 & S4 & S8 & S12 \\\\\n",
    "S1 & S5 & S9 & S13 \\\\\n",
    "S2 & S6 & S10 & S14 \\\\\n",
    "S3 & S7 & S11 & S15\n",
    "\\end{array} \\right)\n",
    "$$\n",
    "\n",
    "into this one:\n",
    "\n",
    "$$\n",
    "state = \\left( \\begin{array}\n",
    "& S0 & S4 & S8 & S12 \\\\\n",
    "S5 & S9 & S13 & S1 \\\\\n",
    "S10 & S14 & S2 & S6 \\\\\n",
    "S15 & S3 & S7 & S11\n",
    "\\end{array} \\right)\n",
    "$$\n",
    "\n",
    "Except for the top row, for which our guess succeeded, the bytes don't align anymore. This isn't technically a difference that should affect anything in our attack. Analyzer, however, only considers one byte key guess at a time and assumes that all the others are zero. This means the leakage was always calculated for a key byte of 0. We could fix this in the leakage model, but really there's no reason to: shift rows is just moving values around, it's not changing anything. This means we don't really learn anything new as compared to the SBox output, at least for a software AES attack. This also means that our regular leakage model will still pick up the SBox operation - the keys will just be shifted from what we would expect.\n",
    "\n",
    "A similar problem affects the next operation in AES, mix columns, with a similarly simple fix. This is covered in our article on extending AES128 attacks to AES256: https://wiki.newae.com/Extending_AES-128_Attacks_to_AES-256. Basically, since Mix Columns is a linear operation, we can continue to use our SBox output. The key we recover, however, will have a shift rows and mix columns operation applied to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.mix_columns_output\n",
    "attack = cwa.cpa(project, leak_model)\n",
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attacking from the other end\n",
    "\n",
    "As we've seen, there's not much point to advancing past the SBox for our software AES attack. What about attacking from the other end? While the other leakage models didn't really any advantages over the SBox output for a software AES attack, there's actually a big reason why we'd like to try an attack from the other end of AES - we'd be able to use the ciphertext instead of the plaintext. While we know both the plaintext and ciphertext for our lab examples, this might not always be the case. For example, what if the target was encrypting messages and sending them to another device? We could monitor the communication lines to learn the ciphertext, but learning the plaintext isn't quite as simple.\n",
    "\n",
    "Looking at a block diagram of AES, we can see that the last block is actually different from the rest of the blocks:\n",
    "\n",
    "![](images/AES_Encryption.png)\n",
    "\n",
    "The first thing to notice is that there's no MixColumns operation for the last round, which makes it a little easier to attack than across the first round. We can basically just repeat our SBox attack and recover the key, as can be seen with the LastroundHW leakage model:\n",
    "\n",
    "```python\n",
    "class LastroundHW(AESLeakageHelper):\n",
    "    name = 'HW: AES Last-Round State'\n",
    "    def leakage(self, pt, ct, key, bnum):\n",
    "        # HD Leakage of AES State between 9th and 10th Round\n",
    "        # Used to break SASEBO-GII / SAKURA-G\n",
    "        st10 = ct[self.INVSHIFT_undo[bnum]]\n",
    "        st9 = inv_sbox(ct[bnum] ^ key[bnum])\n",
    "        return st9\n",
    "\n",
    "    def process_known_key(self, inpkey):\n",
    "        return key_schedule_rounds(inpkey, 0, 10)\n",
    "```\n",
    "\n",
    "Let's see if we can get this model to work. Capture some traces near the end of the AES operation (remember we're attacking at the last round),"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and we can repeat the attack with the new leakage model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.last_round_state\n",
    "attack = cwa.cpa(project, leak_model)\n",
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see that the attack is really confident about a key, but it doesn't seem to be the right one. The block diagram shown above doesn't tell the whole story. The AddRoundKey operation doesn't keep adding the same key. Instead, a transformation is applied to the key for each round and that is added to the state instead. The leakage model includes a method called `key_schedule_rounds()` that we can use to transform our last round key into the first round key that we would recover normally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_key = leak_model.key_schedule_rounds(results.key_guess(), 10, 0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
